import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import warnings

warnings.filterwarnings('ignore')
import tensorflow as tf

tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization, Activation, AveragePooling2D, GlobalAveragePooling2D

(x_train,y_train),(x_test,y_test)=tf.keras.datasets.cifar10.load_data()
x_train=x_train.reshape([-1,32,32,3])/255
x_test=x_test.reshape([-1,32,32,3])/255

#因为每一个卷积层都带有批标准化和激活函数relu，所以这里我们用函数将他们合并在一起
def conv2d(x, filters, num_row, num_col, strides=(1, 1)):
    x = Conv2D(filters, (num_row, num_col), strides=strides, padding = 'same')(x)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    return x
#首先对输入数据进行处理
inputs = tf.keras.Input([32,32,3])
x = conv2d(inputs,16,3,3,strides=(1,1))# out:（32，32，16）

# block1：  out:（16，16，128）
branch_1 = conv2d(x,32,1,1,strides=(2,2))
branch_2 = conv2d(x,32,1,1,strides=(2,2))
branch_2 = conv2d(branch_2,32,3,3)
branch_3 = conv2d(x,32,1,1,strides=(2,2))
branch_3 = conv2d(branch_3,32,5,5)
branch_4 = AveragePooling2D((3,3),strides=(1,1),padding='same')(x)
branch_4 = conv2d(branch_4,32,1,1,strides=(2,2))
#将四个分支叠加在一起，共32+32+32+32=128
x = tf.keras.layers.concatenate([branch_1,branch_2,branch_3,branch_4],axis=3)

# block2：  out:（16，16，128）
branch_1 = conv2d(x,32,1,1)
branch_2 = conv2d(x,32,1,1)
branch_2 = conv2d(branch_2,32,3,3)
branch_3 = conv2d(x,32,1,1)
branch_3 = conv2d(branch_3,32,5,5)
branch_4 = AveragePooling2D((3,3),strides=(1,1),padding='same')(x)
branch_4 = conv2d(branch_4,32,1,1)
#将四个分支叠加在一起，共32+32+32+32=128
x = tf.keras.layers.concatenate([branch_1,branch_2,branch_3,branch_4],axis=3)

# block3：  out:（8，8，256）
branch_1 = conv2d(x,64,1,1,strides=(2,2))
branch_2 = conv2d(x,64,1,1,strides=(2,2))
branch_2 = conv2d(branch_2,64,3,3)
branch_3 = conv2d(x,64,1,1,strides=(2,2))
branch_3 = conv2d(branch_3,64,5,5)
branch_4 = AveragePooling2D((3,3),strides=(1,1),padding='same')(x)
branch_4 = conv2d(branch_4,64,1,1,strides=(2,2))
#将四个分支叠加在一起，共64*4=256
x = tf.keras.layers.concatenate([branch_1,branch_2,branch_3,branch_4],axis=3)

# block4：  out:（8，8，256）
branch_1 = conv2d(x,64,1,1)
branch_2 = conv2d(x,64,1,1)
branch_2 = conv2d(branch_2,64,3,3)
branch_3 = conv2d(x,64,1,1)
branch_3 = conv2d(branch_3,64,5,5)
branch_4 = AveragePooling2D((3,3),strides=(1,1),padding='same')(x)
branch_4 = conv2d(branch_4,64,1,1)
#将四个分支叠加在一起，共64*4=256
x = tf.keras.layers.concatenate([branch_1,branch_2,branch_3,branch_4],axis=3)

# 平均池化后全连接。
x = GlobalAveragePooling2D(name='avg_pool')(x)# (256)
x = Dense(10, activation='softmax', name='predictions')(x)

model = tf.keras.Model(inputs,x)

model.summary()

model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              optimizer=tf.keras.optimizers.Adam(lr=0.001),
              metrics=['accuracy'])

history=model.fit(x_train,y_train,batch_size=64,epochs=1,validation_split=0.3)
